package com.phoenix.shuaidatabase.utils;

import java.util.concurrent.ConcurrentHashMap;

public class LuaRateLimiter {

    private ConcurrentHashMap<String,TokenBucket> sources;

    private static class TokenBucket {

        private String script =
                "-- 获取令牌的数量\n" +
                "local tokens = tonumber(redis.call('get', KEYS[1]))\n" +
                "-- 令牌桶容量\n" +
                "local capacity = tonumber(ARGV[1])\n" +
                "-- 令牌桶填充速率\n" +
                "local rate = tonumber(ARGV[2])\n" +
                "-- 时间间隔\n" +
                "local interval = tonumber(ARGV[3])\n" +
                "\n" +
                "-- 如果令牌桶为空，需要等待\n" +
                "if tokens == nil then\n" +
                "    tokens = capacity - 1\n" +
                "    redis.call('set', KEYS[1], tokens, 'PX', interval)\n" +
                "    return 1\n" +
                "end\n" +
                "\n" +
                "-- 如果令牌桶未满，则生成新的令牌\n" +
                "if tokens < capacity then\n" +
                "    local now = tonumber(redis.call('time')[1])\n" +
                "    local last_time = tonumber(redis.call('get', KEYS[2]) or 0)\n" +
                "    local delta = math.floor((now - last_time) * rate)\n" +
                "    tokens = tokens + delta\n" +
                "    if tokens > capacity then\n" +
                "        tokens = capacity\n" +
                "    end\n" +
                "    redis.call('set', KEYS[1], tokens, 'PX', interval)\n" +
                "    redis.call('set', KEYS[2], now, 'PX', interval)\n" +
                "end\n" +
                "\n" +
                "-- 如果令牌桶已满，则拒绝请求\n" +
                "if tokens < 1 then\n" +
                "    return 0\n" +
                "end\n" +
                "\n" +
                "-- 消耗一个令牌，并返回成功\n" +
                "tokens = tokens - 1\n" +
                "redis.call('set', KEYS[1], tokens, 'PX', interval)\n" +
                "return 1\n";
    }


}
